import torch
from torch_geometric.loader import DataLoader
from dataset import my_dataset
from model import Net
import os


def create_error_data(batch):
    # print(batch)
    batch = list(zip(*batch))
    return tuple(batch)


MODEL_PATH = "/home/Dyf/code/models/gnn/"

model = Net()
model = model.to("cuda")
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = \
    torch.nn.BCELoss()
# print(my_dataset, "my_dataset")
train_loader = DataLoader(my_dataset, batch_size=2,shuffle=True, collate_fn=create_error_data)
best_loss = 10


def train():
    model.train()
    loss_all = 0
    for data in train_loader:
        # print(len(data))
        # print(type(data[0]),len(data))
        data = data.to("cuda")
        optimizer.zero_grad()
        ouput = model(data)
        label = data.y
        loss = criterion(ouput, label)
        loss.backward()
        loss_all += data.num_graphs * loss.item()
        optimizer.step()
    return loss_all / len(my_dataset)


def epoch_train():
    global best_loss
    for epoch in range(401):
        loss = train()
        # print("loss",loss)
        if loss < best_loss:
            best_loss = loss
            torch.save(model, os.path.join(MODEL_PATH, 'Net_{}.pth').format(epoch))
            print("=> saved best model", epoch,loss)
        # if epoch % 10 == 0:
        #     print(loss)


# 使用 train_mask 来确定34个节点中只有4个分类标记，就可以将任务执行的非常好，
# 很好的验证了半监督能力。
epoch_train()
